{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CRIC Convergence Generation\n",
    "\n",
    "The notebook generates and saves everything for plotting: \"kidney/convergence_plots/Plots.ipynb\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data and train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xgboost\n",
    "import pickle\n",
    "import shap\n",
    "from time import time as t\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler, Imputer\n",
    "import sklearn\n",
    "import matplotlib.pyplot as pl\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "X = pd.read_csv(\"data/CRIC_time_4yearESRD_X.csv\")\n",
    "y = np.loadtxt(\"data/CRIC_time_4yearESRD_y.csv\")\n",
    "\n",
    "name_map = {\n",
    "    \"CREATININE_SERUM\": \"Blood Creatinine (determines eGFR)\",\n",
    "    \"PCR_URINE_COMBINED\": \"Urine Protein/Creatinine Ratio\",\n",
    "    \"CYSC_CALIBRATED\": \"Blood Cystatin-C\",\n",
    "    \"UPROTEIN24H\": \"Urine Protein\",\n",
    "    \"SERUM_UREA_NITROGEN\": \"Blood Urea Nitrogen\", # (mg/dL)\n",
    "    \"CRCL\": \"Creatinine clearance\",\n",
    "    \"SERUM_ALBUMIN\": \"Blood albumin\",\n",
    "    \"FGF23_CALIBRATED\": \"Fibroblast growth factor 23\",\n",
    "    \"SYSTOLIC\": \"Systolic blood pressure\",\n",
    "    \"ANEMIA\": \"Anemia\",\n",
    "    \"STOFHLA_SCORE\": \"Health literacy score\",\n",
    "    \"CBCHEMOGLOBIN\": \"Blood Hemoglobin\",\n",
    "    \"HEMATOCRIT\": \"Hematocrit\",\n",
    "    \"PULSE_PRESSURE\": \"Pulse pressure\",\n",
    "    \"WBC\": \"White blood cells\",\n",
    "    \"BIRTH_YEAR\": \"Year of birth\"\n",
    "}\n",
    "mapped_feature_names = list(map(lambda x: name_map.get(x, x), X.columns))\n",
    "\n",
    "# split by patient id\n",
    "pids = np.unique(X[\"PID\"].values)\n",
    "train_pids,test_pids = train_test_split(pids, random_state=0)\n",
    "strain_pids,valid_pids = train_test_split(train_pids, random_state=0)\n",
    "\n",
    "# find the indexes of the samples from the patient ids\n",
    "train_inds = np.where([p in train_pids for p in X[\"PID\"].values])[0]\n",
    "strain_inds = np.where([p in strain_pids for p in X[\"PID\"].values])[0]\n",
    "valid_inds = np.where([p in valid_pids for p in X[\"PID\"].values])[0]\n",
    "test_inds = np.where([p in test_pids for p in X[\"PID\"].values])[0]\n",
    "\n",
    "# create the split datasets\n",
    "X = X.drop([\"PID\"], axis=1)\n",
    "X_train = X.iloc[train_inds,:]\n",
    "X_strain = X.iloc[strain_inds,:]\n",
    "X_valid = X.iloc[valid_inds,:]\n",
    "X_test = X.iloc[test_inds,:]\n",
    "y_train = y[train_inds]\n",
    "y_strain = y[strain_inds]\n",
    "y_valid = y[valid_inds]\n",
    "y_test = y[test_inds]\n",
    "\n",
    "# mean impute for linear and deep models\n",
    "imp = Imputer()\n",
    "imp.fit(X_strain)\n",
    "X_strain_imp = imp.transform(X_strain)\n",
    "X_train_imp = imp.transform(X_train)\n",
    "X_valid_imp = imp.transform(X_valid)\n",
    "X_test_imp = imp.transform(X_test)\n",
    "X_imp = imp.transform(X)\n",
    "\n",
    "# standardize\n",
    "scaler = StandardScaler()\n",
    "scaler.fit(X_strain_imp)\n",
    "X_strain_imp = scaler.transform(X_strain_imp)\n",
    "X_train_imp = scaler.transform(X_train_imp)\n",
    "X_valid_imp = scaler.transform(X_valid_imp)\n",
    "X_test_imp = scaler.transform(X_test_imp)\n",
    "X_imp = scaler.transform(X_imp)\n",
    "\n",
    "# Convert to numpy\n",
    "X_strain_values = X_strain.values\n",
    "X_valid_values = X_valid.values\n",
    "\n",
    "params = {\n",
    "    \"colsample_bytree\": 0.15,\n",
    "    \"max_depth\": 5,\n",
    "    \"reg_alpha\": 0,\n",
    "    \"reg_lambda\": 0,\n",
    "    \"subsample\": 0.7\n",
    "}\n",
    "xgb_model = xgboost.XGBClassifier(\n",
    "    max_depth=params[\"max_depth\"],\n",
    "    n_estimators=10000,\n",
    "    learning_rate=0.005,#math.pow(10, params[\"learning_rate\"]),\n",
    "    subsample=params[\"subsample\"],\n",
    "    reg_lambda=params[\"reg_lambda\"],\n",
    "    colsample_bytree=params[\"colsample_bytree\"],\n",
    "    reg_alpha=params[\"reg_lambda\"],\n",
    "    n_jobs=16,\n",
    "    random_state=1\n",
    ")\n",
    "xgb_model.fit(\n",
    "    X_strain_values, y_strain, verbose=100,\n",
    "    eval_set=[(X_valid_values, y_valid)],\n",
    "    eval_metric=\"logloss\",\n",
    "    early_stopping_rounds=200\n",
    ")\n",
    "min_valid_loss = np.min(xgb_model.evals_result_[\"validation_0\"][\"logloss\"])\n",
    "pickle.dump(xgb_model, open(\"xgb_model_final.pickle.dat\", \"wb\"))\n",
    "\n",
    "xgb_model = pickle.load(open(\"xgb_model_final.pickle.dat\", \"rb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get tree shap explanations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Parallelize across samples and save to files - then compose into a single numpy array\n",
    "import multiprocessing\n",
    "import time\n",
    "import random\n",
    "\n",
    "def explain(ind,N):\n",
    "    lb = int((X_strain.values.shape[0]/N) * ind)\n",
    "    ub = int((X_strain.values.shape[0]/N) * (ind+1))\n",
    "    print(\"Working on interval:\",lb,ub)\n",
    "    indep_expl = shap.TreeExplainer(xgb_model, X_strain.values, feature_dependence=\"independent\")\n",
    "    phi = indep_expl.shap_values(X_strain.values[lb:ub,:])\n",
    "    np.save(\"scratch/CRIC_indep_shap_values_{}.npy\".format(ind), phi)\n",
    "\n",
    "def save_indep_tree_shap_values(N=50):\n",
    "    processes = []\n",
    "    for i in range(N):\n",
    "        proc = multiprocessing.Process(target=explain, args=(i,N,))\n",
    "        processes.append(proc)\n",
    "        proc.start()\n",
    "\n",
    "    for one_process in processes:\n",
    "        one_process.join()\n",
    "\n",
    "    print(\"Done!\")\n",
    "\n",
    "    CRIC_indep_shap_values = []\n",
    "    for i in range(N):\n",
    "        CRIC_indep_shap_values.append(np.load(\"scratch/CRIC_indep_shap_values_{}.npy\".format(i)))\n",
    "    np.save(\"CRIC_indep_shap_values_final.npy\", np.concatenate(CRIC_indep_shap_values,0))\n",
    "\n",
    "save_indep_tree_shap_values()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CRIC_indep_shap_values = np.load(\"CRIC_indep_shap_values_final.npy\")\n",
    "print(\"Independent Tree SHAP - top 10\\n\"+\"-\"*20)\n",
    "print(np.array(list(X_strain.columns.values))[np.argsort(-np.abs(CRIC_indep_shap_values.sum(0)))[0:10]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 10th biggest\n",
    "one_cutoff = np.sort(np.mean(np.abs(CRIC_indep_shap_values),0))[-10]*.01\n",
    "fiv_cutoff = np.sort(np.mean(np.abs(CRIC_indep_shap_values),0))[-10]*.05\n",
    "ten_cutoff = np.sort(np.mean(np.abs(CRIC_indep_shap_values),0))[-10]*.1\n",
    "print(\"10% Cutoff (10% of tenth largest global shap value):\",round(ten_cutoff,6))\n",
    "print(\"5% Cutoff (5% of tenth largest global shap value)  :\",round(fiv_cutoff,6))\n",
    "print(\"1% Cutoff (1% of tenth largest global shap value)  :\",round(one_cutoff,6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(20)\n",
    "X_values = X_strain_values\n",
    "X_values[np.isnan(X_values)] = 0\n",
    "\n",
    "x_strain_mean = np.nanmean(X_strain_values,0).reshape([1,X_strain_values.shape[1]])\n",
    "mismatch_inds = np.array(range(0,6064))[xgb_model.predict(X_strain_values)!=xgb_model.predict(x_strain_mean)]\n",
    "x_inds = np.random.choice(mismatch_inds, 10, replace = False)\n",
    "np.argsort(xgb_model.predict(X_strain_values)!=xgb_model.predict(x_strain_mean))[-11:-1]\n",
    "ten_x = X_strain_values[x_inds,:]\n",
    "# Make sure all ten differ to the mean in terms of the model prediction\n",
    "assert np.sum(xgb_model.predict(x_strain_mean)[0] != xgb_model.predict(ten_x)) == 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run and save sampling based results (kernel and IME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "## Sample Explainer\n",
    "for x_ind in x_inds:\n",
    "    ns_lst = np.power(10,[i/2 for i in range(3*2,7*2)])\n",
    "    num_iter = 10\n",
    "    \n",
    "    print(\"x_ind\",x_ind)\n",
    "    directory = \"samplev2_{}/\".format(x_ind)\n",
    "    if not os.path.exists(directory): os.makedirs(directory)\n",
    "    \n",
    "    f = lambda inp : xgb_model.predict(inp, output_margin=True)\n",
    "    se = shap.SamplingExplainer(f, x_strain_mean)\n",
    "    ke = shap.KernelExplainer(f, x_strain_mean)\n",
    "    x = X_strain_values[x_ind:x_ind+1,:]\n",
    "    \n",
    "    expl_tree = shap.TreeExplainer(xgb_model, x_strain_mean, feature_dependence=\"independent\", model_output=\"margin\")\n",
    "    itshap = expl_tree.shap_values(x)\n",
    "    \n",
    "    for ns in ns_lst:\n",
    "        mspf_lst = [10, 50, 100, 500, 1000, ns]\n",
    "        for mspf in mspf_lst:\n",
    "            # Make sure number of samples is divisible by 2\n",
    "            ns = int(ns)\n",
    "            if (int(ns)%2 == 1): ns = ns+1\n",
    "\n",
    "            print(\"ns\",ns,\"mspf\",mspf)\n",
    "\n",
    "            # Time how long it takes to obtain shap_values\n",
    "            s = t()\n",
    "            se_out = np.vstack([se.shap_values(x, silent=True, nsamples=ns, \n",
    "                                               min_samples_per_feature=int(mspf)) for i in range(num_iter)])\n",
    "            runtime = (t() - s)/float(num_iter)\n",
    "            assert np.abs(se_out[0].sum() - itshap.sum()) < .01\n",
    "            np.save(directory+\"se_out_ns:{}_mspf:{}_runtime:{}\".format(ns,mspf,runtime), se_out)\n",
    "\n",
    "    l1reg_lst = [0.0, \"aic\", \"bic\"]\n",
    "    for ns in ns_lst:\n",
    "        for l1reg in l1reg_lst:\n",
    "            # Make sure number of samples is divisible by 2\n",
    "            ns = int(ns)\n",
    "            if (int(ns)%2 == 1): ns = ns+1\n",
    "        \n",
    "            print(\"ns\",ns,\"l1reg\",l1reg)\n",
    "\n",
    "            # Time how long it takes to obtain shap_values\n",
    "            s = t()\n",
    "            ke_out = np.vstack([ke.shap_values(x, silent=True, nsamples=ns, \n",
    "                                               l1_reg=l1reg) for i in range(num_iter)])\n",
    "            runtime = (t() - s)/float(num_iter)\n",
    "            assert np.abs(ke_out[0].sum() - itshap.sum()) < .01\n",
    "            np.save(directory+\"ke_out_ns:{}_l1reg:{}_runtime:{}\".format(ns,l1reg,runtime), ke_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load and compute deviations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "se_abs_dev_from_truth_mean_single_dict = {}\n",
    "se_abs_dev_from_truth_max_single_dict = {}\n",
    "indep_expl_single_ref = shap.TreeExplainer(xgb_model, x_strain_mean, feature_dependence=\"independent\")\n",
    "for x_ind in x_inds:    \n",
    "    single_phi_x = indep_expl_single_ref.shap_values(X_strain.values[x_ind:x_ind+1,:])\n",
    "    phi_x = CRIC_indep_shap_values[x_ind:x_ind+1,:]\n",
    "    directory = \"samplev2_{}/\".format(x_ind)\n",
    "    ns_lst_pwrs = [i/2 for i in range(3*2,7*2)]\n",
    "    ns_lst = np.power(10,ns_lst_pwrs)\n",
    "    se_out_files = [f for f in os.listdir(directory) if \"se_out\" in f]\n",
    "    for mspf in [10, 50, 100, 500, 1000, \"ns\"]:\n",
    "        abs_dev_from_truth_single_mean_lst = []\n",
    "        abs_dev_from_truth_single_max_lst = []\n",
    "        for ns in ns_lst:\n",
    "            if mspf == \"ns\":\n",
    "                se_out_file = [f for f in se_out_files if \"ns:{}_\".format(r2(ns)) in f \n",
    "                               and \"mspf:{}_\".format(ns) in f]\n",
    "            else:\n",
    "                se_out_file = [f for f in se_out_files if \"ns:{}_\".format(r2(ns)) in f \n",
    "                               and \"mspf:{}_\".format(mspf) in f]\n",
    "            if not se_out_file: continue\n",
    "            se_out_file = se_out_file[0]\n",
    "            se_out = np.load(directory+se_out_file)\n",
    "            assert np.abs(se_out[0].sum() - single_phi_x.sum()) < 0.01\n",
    "            \n",
    "            abs_dev_from_truth_single = np.abs(se_out-single_phi_x).mean(0)\n",
    "            abs_dev_from_truth_single_mean_lst.append(abs_dev_from_truth_single.mean())\n",
    "            abs_dev_from_truth_single_max_lst.append(abs_dev_from_truth_single.max())            \n",
    "\n",
    "        if (mspf not in se_dev_from_truth_mean_dict):\n",
    "            se_abs_dev_from_truth_mean_single_dict[mspf] = np.array(abs_dev_from_truth_single_mean_lst)/len(x_inds)\n",
    "            se_abs_dev_from_truth_max_single_dict[mspf] = np.array(abs_dev_from_truth_single_max_lst)/len(x_inds)\n",
    "        else:\n",
    "            se_abs_dev_from_truth_mean_single_dict[mspf] += np.array(abs_dev_from_truth_single_mean_lst)/len(x_inds)\n",
    "            se_abs_dev_from_truth_max_single_dict[mspf] += np.array(abs_dev_from_truth_single_max_lst)/len(x_inds)\n",
    "\n",
    "sample_expl_abs_dev_single_ref_truth_mean_dict = se_abs_dev_from_truth_mean_single_dict\n",
    "sample_expl_abs_dev_single_ref_truth_max_dict = se_abs_dev_from_truth_max_single_dict\n",
    "pickle.dump( sample_expl_abs_dev_single_ref_truth_mean_dict, open( \"sample_expl_abs_dev_single_ref_truth_mean_dict.p\", \"wb\" ) )\n",
    "pickle.dump( sample_expl_abs_dev_single_ref_truth_max_dict, open( \"sample_expl_abs_dev_single_ref_truth_max_dict.p\", \"wb\" ) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ke_abs_dev_from_truth_mean_single_dict = {}\n",
    "ke_abs_dev_from_truth_max_single_dict = {}\n",
    "l1reg_lst = [0.0, \"aic\", \"bic\"]\n",
    "\n",
    "indep_expl_single_ref = shap.TreeExplainer(xgb_model, x_strain_mean, feature_dependence=\"independent\")\n",
    "for x_ind in x_inds:\n",
    "    single_phi_x = indep_expl_single_ref.shap_values(X_strain.values[x_ind:x_ind+1,:])\n",
    "    phi_x = CRIC_indep_shap_values[x_ind:x_ind+1,:]\n",
    "    directory = \"samplev2_{}/\".format(x_ind)\n",
    "    print(directory)\n",
    "    ns_lst_pwrs = [i/2 for i in range(3*2,7*2)]\n",
    "    ns_lst = np.power(10,ns_lst_pwrs)\n",
    "    ke_out_files = [f for f in os.listdir(directory) if \"ke_out\" in f]\n",
    "    \n",
    "    for l1reg in l1reg_lst:\n",
    "        abs_dev_from_truth_single_mean_lst = []\n",
    "        abs_dev_from_truth_single_max_lst = []\n",
    "        for ns in ns_lst:\n",
    "            ke_out_file = [f for f in ke_out_files if \"ns:{}_\".format(r2(ns)) in f \n",
    "                           and \"l1reg:{}_\".format(l1reg) in f]\n",
    "            if not ke_out_file: continue\n",
    "            ke_out_file = ke_out_file[0]\n",
    "            ke_out = np.load(directory+ke_out_file)\n",
    "            assert np.abs(ke_out[0].sum() - single_phi_x.sum()) < 0.01\n",
    "            \n",
    "            abs_dev_from_truth_single = np.abs(ke_out-single_phi_x).mean(0)\n",
    "            abs_dev_from_truth_single_mean_lst.append(abs_dev_from_truth_single.mean())\n",
    "            abs_dev_from_truth_single_max_lst.append(abs_dev_from_truth_single.max())                        \n",
    "            \n",
    "        if (l1reg == 0.0): print(max_std_lst)\n",
    "            \n",
    "        if (l1reg not in ke_dev_from_truth_mean_dict):\n",
    "            ke_abs_dev_from_truth_mean_single_dict[l1reg] = np.array(abs_dev_from_truth_single_mean_lst)/len(x_inds)\n",
    "            ke_abs_dev_from_truth_max_single_dict[l1reg] = np.array(abs_dev_from_truth_single_max_lst)/len(x_inds)\n",
    "        else:\n",
    "            ke_abs_dev_from_truth_mean_single_dict[l1reg] += np.array(abs_dev_from_truth_single_mean_lst)/len(x_inds)\n",
    "            ke_abs_dev_from_truth_max_single_dict[l1reg] += np.array(abs_dev_from_truth_single_max_lst)/len(x_inds)\n",
    "\n",
    "kernel_expl_abs_dev_single_ref_truth_mean_dict = ke_abs_dev_from_truth_mean_single_dict\n",
    "kernel_expl_abs_dev_single_ref_truth_max_dict = ke_abs_dev_from_truth_max_single_dict\n",
    "pickle.dump( kernel_expl_abs_dev_single_ref_truth_mean_dict, open( \"kernel_expl_abs_dev_single_ref_truth_mean_dict.p\", \"wb\" ) )\n",
    "pickle.dump( kernel_expl_abs_dev_single_ref_truth_max_dict, open( \"kernel_expl_abs_dev_single_ref_truth_max_dict.p\", \"wb\" ) )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
